"""
ConnLLM Moonshot网络搜索能力V2实现
提供Moonshot API网络搜索功能
"""

from typing import Dict, Any, List, Optional, Union
import json
import logging
import uuid

from ...core_v2.capability import Capability

logger = logging.getLogger("connllm.implementations.moonshot.web_search_v2")

class MoonshotWebSearchV2(Capability):
	"""Moonshot网络搜索能力V2实现"""
	
	def __init__(self, provider):
		"""
		初始化能力
		
		Args:
			provider: 提供商实例
		"""
		super().__init__(provider)
		self.search_tools = self._prepare_search_tools()
	
	def prepare_request(self, params: Dict[str, Any]) -> Dict[str, Any]:
		"""
		准备请求参数，添加搜索工具
		
		Args:
			params: 请求参数
			
		Returns:
			处理后的请求参数
		"""
		# 如果没有搜索工具，直接返回
		if not self.search_tools:
			return params
		
		# 从params获取当前工具列表，如果没有则创建空列表
		tools = params.get("tools", [])
		
		# 添加搜索工具
		for tool in self.search_tools:
			# 检查工具是否已存在
			exists = False
			for existing_tool in tools:
				if (existing_tool.get("type") == "function" and
					existing_tool.get("function", {}).get("name") == tool.get("function", {}).get("name")):
					exists = True
					break
			
			# 如果不存在，添加到列表
			if not exists:
				tools.append(tool)
		
		# 更新params
		if tools:
			params["tools"] = tools
			
			# 如果未设置tool_choice，默认为auto
			if "tool_choice" not in params:
				params["tool_choice"] = "auto"
		
		return params
	
	def detect_search_calls(self, response: Dict[str, Any]) -> List[Dict[str, Any]]:
		"""
		从响应中检测搜索调用
		
		Args:
			response: API响应
			
		Returns:
			检测到的搜索调用列表
		"""
		search_calls = []
		
		try:
			# 检查响应中是否有工具调用
			if "choices" in response and len(response["choices"]) > 0:
				choice = response["choices"][0]
				
				if "message" in choice and "tool_calls" in choice["message"]:
					raw_tool_calls = choice["message"]["tool_calls"]
					
					# 处理每个工具调用
					for tool_call in raw_tool_calls:
						# 确保是函数调用
						if tool_call.get("type") == "function":
							function = tool_call.get("function", {})
							name = function.get("name", "")
							
							# 检查是否是搜索工具
							if name == "web_search":
								# 提取搜索调用信息
								call_info = {
									"id": tool_call.get("id", str(uuid.uuid4())),
									"name": name,
									"arguments": function.get("arguments", "{}")
								}
								
								# 如果arguments是字符串，尝试解析为JSON
								if isinstance(call_info["arguments"], str):
									try:
										call_info["arguments"] = json.loads(call_info["arguments"])
									except json.JSONDecodeError:
										logger.warning(f"解析搜索参数失败: {call_info['arguments']}")
										call_info["arguments"] = {}
								
								search_calls.append(call_info)
		except Exception as e:
			logger.error(f"检测搜索调用失败: {str(e)}")
		
		return search_calls
	
	def execute_search_calls(self, search_calls: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
		"""
		执行搜索调用
		
		Args:
			search_calls: 搜索调用列表
			
		Returns:
			搜索结果消息列表
		"""
		search_results = []
		
		for call in search_calls:
			call_id = call.get("id", "")
			arguments = call.get("arguments", {})
			
			# 提取搜索查询
			query = arguments.get("query", "")
			if not query:
				logger.warning("搜索查询为空")
				continue
				
			# 执行搜索逻辑
			try:
				# 这里通常会调用搜索API，但为简化实现，我们创建模拟结果
				result = self._perform_search(query)
				
				# 创建搜索结果消息
				if isinstance(result, dict):
					content = json.dumps(result, ensure_ascii=False)
				elif isinstance(result, list):
					content = json.dumps(result, ensure_ascii=False)
				else:
					content = str(result)
				
				search_result = {
					"role": "tool",
					"tool_call_id": call_id,
					"name": "web_search",
					"content": content
				}
				
				search_results.append(search_result)
				
			except Exception as e:
				logger.error(f"执行搜索失败: {str(e)}")
				
				# 创建错误消息
				error_result = {
					"role": "tool",
					"tool_call_id": call_id,
					"name": "web_search",
					"content": json.dumps({"error": f"搜索错误: {str(e)}"}, ensure_ascii=False)
				}
				
				search_results.append(error_result)
		
		return search_results
	
	def _prepare_search_tools(self) -> List[Dict[str, Any]]:
		"""
		准备搜索工具定义
		
		Returns:
			搜索工具定义列表
		"""
		# 创建搜索工具定义
		search_tool = {
			"type": "function",
			"function": {
				"name": "web_search",
				"description": "搜索互联网以获取实时信息，返回搜索结果列表",
				"parameters": {
					"type": "object",
					"properties": {
						"query": {
							"type": "string",
							"description": "搜索查询，应简短明确"
						}
					},
					"required": ["query"]
				}
			}
		}
		
		return [search_tool]
	
	def _perform_search(self, query: str) -> Union[List[Dict[str, Any]], Dict[str, Any]]:
		"""
		执行搜索
		
		Args:
			query: 搜索查询
			
		Returns:
			搜索结果
		"""
		# 从配置中获取搜索设置
		search_config = self.provider.config_manager.get("search_config", {})
		provider = search_config.get("provider", "mock")
		
		# 检查是否有真实搜索提供商配置
		if provider != "mock" and "api_key" in search_config:
			# 这里可以集成真实搜索API，如Google、Bing等
			# 但为了演示，我们使用模拟结果
			logger.info(f"使用搜索提供商: {provider}")
		
		# 生成模拟搜索结果
		return {
			"query": query,
			"results": [
				{
					"title": f"关于'{query}'的搜索结果1",
					"url": f"https://example.com/result1",
					"snippet": f"这是关于'{query}'的第一个搜索结果的摘要。包含一些相关信息和上下文。"
				},
				{
					"title": f"关于'{query}'的搜索结果2",
					"url": f"https://example.com/result2",
					"snippet": f"这是关于'{query}'的第二个搜索结果的摘要。提供了不同的视角和信息。"
				},
				{
					"title": f"关于'{query}'的搜索结果3",
					"url": f"https://example.com/result3",
					"snippet": f"这是关于'{query}'的第三个搜索结果的摘要。包含额外的信息和细节。"
				}
			]
		}
